#include "faiss_index.h"
#include <faiss/IndexIDMap.h>
#include "logger.h"

bool RoaringBitmapIdSelector::is_member(int64_t id)const
{
    return roaring_bitmap_contains(bitmap,static_cast<uint32_t>(id));
}

FaissIndex::FaissIndex(faiss::Index* index) : index(index) {}
std::pair<std::vector<long>,std::vector<float>> FaissIndex::searchVector(const std::vector
<float> &data,int k,const roaring_bitmap_t* bitmap)
{
    //获取需要查询的向量维度
    int dim=index->d;
    int queryNum=data.size()/dim;

    std::vector<long> indices(queryNum*k);
    std::vector<float> distances(queryNum*k);

    //如果传入bitmap，则使用RoaringBitmapIdSelector初始化faiss的searchparamaters
    faiss::SearchParameters params;
    RoaringBitmapIdSelector selector(bitmap);
    if(bitmap!=nullptr){
        params.sel=&selector;
    }

    index->search(queryNum,data.data(),k,distances.data(),indices.data(),&params);
    GlobalLogger->debug("Retrieved values:");
    for(size_t i=0;i<indices.size();i++){
        if(indices[i]!=-1){
            GlobalLogger->debug("ID: {}, Distance: {}",indices[i],distances[i]);
        }else{
            GlobalLogger->debug("No specific ID found");
        }
    }
    return std::make_pair(indices,distances);
}

void FaissIndex::insert(const std::vector<float> &data,uint64_t label)
{
    long id=static_cast<long>(label);
    index->add_with_ids(1,data.data(),&id);
}

void FaissIndex::remove_vectors(const std::vector<long> &ids){
    faiss::IndexIDMap* index_idmap=dynamic_cast<faiss::IndexIDMap*>(index);
    if(index_idmap){
        faiss::IDSelectorBatch selector(ids.size(),ids.data());
        index_idmap->remove_ids(selector);
    }else{
        throw std::runtime_error("Faiss index is not IndexIDMap");
    }
}

void FaissIndex::saveIndex(const std::string& file_path){
    faiss::write_index(index,file_path.c_str());
}
void FaissIndex::loadIndex(const std::string& file_path){
    std::ifstream file(file_path);
    if(file.good()){
        file.close();
        if(index!=nullptr){
            delete index;
        }
        index=faiss::read_index(file_path.c_str());
    }else{
        GlobalLogger->warn("FaissIndex::loadIndex: file not found");
    }
}